import math
import sys,os,shutil
import heapq
import debugpy
import random
sys.path.append("/home/fd_chen/in_context_RL/spoc-robot-training")

from environment.spoc_objects import SPOCObject
from utils.debug_utils import enable_remote_debug
from utils.constants.stretch_initialization_utils import STRETCH_ENV_ARGS, AGENT_ROTATION_DEG, AGENT_MOVEMENT_CONSTANT
from environment.stretch_controller import StretchController
from matplotlib import pyplot as plt
from path_A_star_expert.actions import determine_turn_actions,determine_turn_actions_end,ACTIONS
from path_A_star_expert.astar_utils import  save_navigation_frame ,get_one_object,save_data_json,filter_objects,target_types,get_position_rotation,get_top_down_map,load_procthor_houses


class Node:
    def __init__(self, x, z, cost, pind, action=None,agent_pose=None,object_position=None):
        self.x = x  # x position of node
        self.z = z  # y position of node
        self.cost = cost  # g cost of node
        self.pind = pind  # parent index of node
        self.action = action  # action to reach this node
        self.agent_pose = agent_pose # agent_pose used to teleport
        self.object_position = object_position # object position used to teleport

class StarExpert: 
    def __init__(self,house_id=0,scene=None):
        self.id = house_id

        STRETCH_ENV_ARGS['gpu_device'] = 0
        # scene 
        self.stretch_controller = StretchController(**STRETCH_ENV_ARGS)
        self.stretch_controller.reset(scene)

        self.reachable_positions = self.stretch_controller.get_reachable_positions()  
        self.bounds = self.stretch_controller.controller.last_event.metadata["sceneBounds"]["size"]# width and height of the scene
        self.motion = self.get_motion()
        self.obsmap_positions = self.obsmap()
        self.successuful_objects = []
        
        
    def reset_task(self, target: SPOCObject):

        self.target = target
        self.goal_node = self.target_goal()
        self.start_node = self.agent_start()
        self.pathx, self.pathz = [], []
        self.shortest_actions = []
        self.shortest_frames = []
        self.shortest_positions = []
        self.early_stop = False
        
        self.build_task_dir()

    def get_actions_from_nodes(self, nodes):
        actions = []
        for node in nodes:
            actions.extend(node.action)
        # actions.append(ACTIONS['end'])
        return actions
    
    def get_nav_frames_positions_from_nodes(self, nodes,visualize=False):
        self.teleport_agent(nodes[0].agent_pose)
        self.shortest_frames,self.shortest_positions= [self.stretch_controller.navigation_camera],[get_position_rotation(self.stretch_controller)] # initial frame and position
        for action in self.shortest_actions[1:]: #　１：start　-1:end
            self.stretch_controller.agent_step(action)
            if not self.stretch_controller.controller.last_event.metadata["lastActionSuccess"]:
                self.teleport_agent(self.start_node.agent_pose)
                raise ValueError("Action failed,agent teleported back")
            self.shortest_frames.append(self.stretch_controller.navigation_camera)
            self.shortest_positions.append(get_position_rotation(self.stretch_controller))
        
        if not self.is_visuable():
            self.teleport_agent(self.start_node.agent_pose)
            raise ValueError("Target object is not visuable,agent teleported back")

        # add the last frame and position the same as the end node

        assert len(self.shortest_frames) == len(self.shortest_actions)
        if visualize: #FIXME: add visualization
            pass
        return self.shortest_frames,self.shortest_positions
    
    def save_task(self):

        self.shortest_actions = self.get_actions_from_nodes(self.shortest_nodes) # make sure the actions are correct add the end action    
        self.shortest_frames,self.shortest_positions = self.get_nav_frames_positions_from_nodes(self.shortest_nodes)
        self.successuful_objects.append(self.target['objectId'])

        save_navigation_frame(self.shortest_frames, self.frame_dir)
        get_top_down_map(self.stretch_controller,self.shortest_positions, save_path=f"{self.metadir}/top_down_map.png")
        save_data_json(self.shortest_positions, f"{self.metadir}/position.json")
        save_data_json(self.shortest_actions, f"{self.metadir}/actions.json")
        save_data_json({**self.target._thor_obj}, f"{self.metadir}/target.json")
        save_data_json(self.successuful_objects,f"{self.root_dir}/successuful_objects.json")
        draw(self, name=f"{self.metadir}/path")
        # FIXME: ADD THE HOUSE SCENE
        
    
    def build_task_dir(self):
        self.root_dir = f"./debug/house_{self.id}"
        if not os.path.exists(self.root_dir):
            os.makedirs(self.root_dir)
        self.frame_dir = f"{self.root_dir}/{self.target['objectType']+'|'+self.target['objectId']}/frames"
        self.metadir = f"{self.root_dir}/{self.target['objectType']+'|'+self.target['objectId']}/meta"
        for dir in [self.frame_dir, self.metadir]:
            if os.path.exists(dir):
                shutil.rmtree(dir)
            os.makedirs(dir, exist_ok=True)

            
    def target_goal(self):
        self.target_params = {
            'x': self.stretch_controller.get_object_position(self.target['objectId'])['x'],
            'z': self.stretch_controller.get_object_position(self.target['objectId'])['z'],
        }
        return Node(self.target_params['x'], self.target_params['z'], 0.0, -1,action='end',object_position=self.stretch_controller.get_object_position(self.target['objectId']))
    
    def agent_start(self):
        self.agent_start_params = {
            'x': self.stretch_controller.get_current_agent_full_pose()['position']['x'],
            'z': self.stretch_controller.get_current_agent_full_pose()['position']['z'],
        }
        return Node(self.agent_start_params['x'], self.agent_start_params['z'], 0.0, -1,action=' ',agent_pose=self.stretch_controller.get_current_agent_full_pose())
    
    def action_from_dx_dz(self,dx, dz,n_curr:Node, teleport_agent_to_pind:bool = True):

        actions = []
        target_angle = round(math.degrees(math.atan2(dx, dz))) % 360
        turn_actions, current_heading = determine_turn_actions(round(n_curr.agent_pose['rotation']['y']), target_angle)
        
        actions.extend(turn_actions)
        if ACTIONS['move_back'] not in turn_actions:
            actions.append(ACTIONS['move_ahead'])
        
        for action in actions:
            self.stretch_controller.agent_step(action)
            if not self.stretch_controller.controller.last_event.metadata["lastActionSuccess"]:
                return None,None,None
        
        node_agent_pose = self.stretch_controller.get_current_agent_full_pose()
        
        if teleport_agent_to_pind:
            self.teleport_agent(n_curr.agent_pose)
            
        return actions,node_agent_pose,current_heading


    def visualize_env(self, mode="unreachable"):
        fig, ax = plt.subplots(1, 1)
        
        xs = [rp["x"] for rp in self.reachable_positions]
        zs = [rp["z"] for rp in self.reachable_positions]
        ax.scatter(xs, zs, color="blue")
        
        if mode == "unreachable":
            xs_obs = [obs["x"] for obs in self.obsmap_positions]
            zs_obs = [obs["z"] for obs in self.obsmap_positions]
            ax.scatter(xs_obs, zs_obs, color="red")
            
        ax.set_xlabel("$x$")
        ax.set_ylabel("$z$")
        ax.set_title("Positions in the Scene")
        ax.set_aspect("equal")
        plt.savefig(f"{self.metadir}/positions_in_the_Scence.png")
        plt.close()

    def is_reachable(self, x, z):
        for rp in self.reachable_positions:
            distance = math.sqrt((rp["x"] - x) ** 2 + (rp["z"] - z) ** 2)
            if distance < 0.05:
                return True
        return False

    def obsmap(self, vis=False):  # sourcery skip: use-itertools-product
        self.obsmap_positions = []
        self.min_x, self.max_x = min(rp["x"] for rp in self.reachable_positions), max(
            rp["x"] for rp in self.reachable_positions
        )
        self.min_z, self.max_z = min(rp["z"] for rp in self.reachable_positions), max(
            rp["z"] for rp in self.reachable_positions
        )

        for x_ind in range(round((self.max_x - self.min_x) / STRETCH_ENV_ARGS['gridSize']) + 1):
            for z_ind in range(round((self.max_z - self.min_z) / STRETCH_ENV_ARGS['gridSize']) + 1):
                x = self.min_x + x_ind * STRETCH_ENV_ARGS['gridSize']
                z = self.min_z + z_ind * STRETCH_ENV_ARGS['gridSize']
                if not self.is_reachable(x, z):
                    self.obsmap_positions.append({"x": x, "z": z})

        if vis:
            self.visualize_env()
        
        return self.obsmap_positions
    
    def teleport_agent(self, agent_pose, teleport_horizon=False):
        """Teleport agent to a given pose."""
        action_params = {
            'action': 'Teleport',
            'position': agent_pose['position'],
            'rotation': agent_pose['rotation'],
            'standing': agent_pose['isStanding']
        }
        if teleport_horizon:
            action_params['horizon'] = agent_pose["cameraHorizon"]
            
        return self.stretch_controller.controller.step(**action_params).metadata["lastActionSuccess"]

    def get_motion(self, min_angle_deg=AGENT_ROTATION_DEG, move_distance=AGENT_MOVEMENT_CONSTANT, small_probility=0.2):

        self.motion = []
        min_angle_rad = math.radians(min_angle_deg)  # Convert to radians
        num_directions = int(360 / min_angle_deg)   # Total number of directions

        for i in range(num_directions):
            theta = i * min_angle_rad  # Radians of the current direction
            dx = round(move_distance * math.cos(theta), 3)  # x direction component, rounded to 3 decimal places
            dy = round(move_distance * math.sin(theta), 3)  # y direction component

            if abs(i * min_angle_deg) % 30 != 0 and random.random() < small_probility:
                self.motion.append([dx, dy])  
            elif abs(i * min_angle_deg) % 30 == 0:
                self.motion.append([dx, dy])  

        return self.motion

    def calc_index(self, node):
        return round((node.z - self.min_z) * self.bounds['x'] + (node.x - self.min_x), 3)

    def fvalue(self, node, n_goal):
        return node.cost + self.h(node, n_goal)
    
    def h(self, node, n_goal):
        return math.hypot(node.x - n_goal.x, node.z - n_goal.z)
    
    def extract_path(self, closed_set):
        self.shortest_nodes = []
        if not self.early_stop:
            pathx, pathz = [self.goal_node.x], [self.goal_node.z]
            n_ind = self.calc_index(self.goal_node)
        else:
            pathx, pathz = [self.end_node.x], [self.end_node.z]
            n_ind = self.calc_index(self.end_node)
            
        while True:
            node = closed_set[n_ind]
            pathx.append(node.x)
            pathz.append(node.z)
            self.shortest_nodes.append(node) # include the start node and but uninclude the end node position
            n_ind = node.pind

            if node == self.start_node:
                break

        pathx = list(reversed(pathx))
        pathz = list(reversed(pathz))
        shortes_nodes = list(reversed(self.shortest_nodes)) # include the start node and end node position
        return pathx, pathz, shortes_nodes


    def a_start_planner(self):
        self.open_set, self.closed_set = {}, {}
        q_priority = []
        self.open_set[self.calc_index(self.start_node)] = self.start_node
        heapq.heappush(q_priority,(self.fvalue(self.start_node, self.goal_node), self.calc_index(self.start_node)))

        while self.open_set:
            _, ind = heapq.heappop(q_priority)
            n_curr = self.open_set[ind]
            self.closed_set[ind] = n_curr
            self.open_set.pop(ind)
            print("open_set", len(self.open_set))

            if self.check_distance(n_curr, self.goal_node, threhold=1.0):
                return self.shortest_nodes

            for i in range(len(self.motion)): # one step 
                
                detla_x, detla_z = self.motion[i]
                node = Node(n_curr.x + detla_x,
                            n_curr.z + detla_z,
                            n_curr.cost + AGENT_MOVEMENT_CONSTANT, ind)# index parent node

                if not self.check_node(node):
                    continue

                actions,agent_pose,_ = self.action_from_dx_dz(detla_x, detla_z, n_curr)
                if actions is None:
                    continue

                n_ind = self.calc_index(node) # 拓展节点的index

                if n_ind not in self.closed_set:
                    if n_ind in self.open_set:
                        if self.open_set[n_ind].cost > node.cost:
                            self.open_set[n_ind].cost = node.cost
                            self.open_set[n_ind].pind = ind
                            self.open_set[n_ind].action = actions
                            self.open_set[n_ind].agent_pose = agent_pose
                    else:
                        node.action = actions # more than one action
                        node.agent_pose = agent_pose
                        self.open_set[n_ind] = node
                        heapq.heappush(q_priority, (self.fvalue(node, self.goal_node), self.calc_index(node)))

    
    def check_distance(self, node1, node2, threhold=1):
        if math.hypot(node1.x - node2.x, node1.z - node2.z) < threhold:
            self.early_stop = True
            self.end_node = node1
            self.pathx, self.pathz , self.shortest_nodes = self.extract_path(self.closed_set)
            self.shortest_path = [[x,z] for x, z in zip(self.pathx,self.pathz)]
            return True
        return False

    def check_node(self, node, threhold=0.25):  # FIXME:add collision check and action check
        
        if node.x <= self.min_x or node.x >= self.max_x or \
                node.z <= self.min_z or node.z >= self.max_z:
            return False

        for obs in self.obsmap_positions:
            distance = math.sqrt((obs["x"] - node.x) ** 2 + (obs["z"] - node.z) ** 2)
            if distance <= threhold:
                return False  # in collision
        return True
    
    def is_visuable(self):
        debugpy.breakpoint()
        if not self.stretch_controller.is_object_visible_enough_for_interaction(self.target['objectId'],manipulation_camera=False):
            angle_diff = self.stretch_controller.get_agent_alignment_to_object(self.target['objectId'])
            end_actions,_ = determine_turn_actions_end(self.stretch_controller.get_current_agent_full_pose()['rotation']['y'],angle_diff)
            for action in end_actions:
                self.stretch_controller.agent_step(action)
                if not self.stretch_controller.controller.last_event.metadata["lastActionSuccess"]:
                    self.teleport_agent(self.start_node.agent_pose)
                    raise ValueError("Action failed,agent teleported back")
                self.shortest_actions.append(action)
                self.shortest_positions.append(get_position_rotation(self.stretch_controller))
                self.shortest_frames.append(self.stretch_controller.navigation_camera)
                if self.stretch_controller.is_object_visible_enough_for_interaction(self.target['objectId'],manipulation_camera=False):
                    break
        return self.stretch_controller.is_object_visible_enough_for_interaction(self.target['objectId'],manipulation_camera=False)

def draw(expert, name="./debug/init.png"):
        fig, ax = plt.subplots(1, 1)
                
        xs_obs = [obs["x"] for obs in expert.obsmap_positions]
        zs_obs = [obs["z"] for obs in expert.obsmap_positions]
        plt.plot(xs_obs, zs_obs, 'sk')
        plt.plot(expert.pathx, expert.pathz, '-r')
        plt.plot(expert.start_node.x, expert.start_node.z, 'sg')
        plt.plot(expert.goal_node.x, expert.goal_node.z, 'sb')
        
        plt.axis("equal")
        plt.show()
        plt.savefig(f"{name}.png")


if __name__ == '__main__':
    # enable_remote_debug(65532)
    houses = load_procthor_houses()
    house_id = 2
    # for house in houses:
    star_expert = StarExpert(house_id,scene=houses[house_id])
    objects = filter_objects(star_expert.stretch_controller,target_types())
    for target in objects:
        try:
            print(f"Target object: {target['objectId']},{target['objectType']}")
            star_expert.reset_task(target)
            star_expert.a_start_planner()
            star_expert.save_task()
        except Exception as e:
            print(e)
            continue
    star_expert.stretch_controller.stop()
    house_id += 1
    
    
